package simpletrace

import (
	"math/rand"
	"sync"
	"time"
)

// Sampler 是采样策略接口
type Sampler interface {
	ShouldSample(traceID string) bool
}

// ProbabilitySampler 实现按概率采样
type ProbabilitySampler struct {
	rate float64 // 采样概率 [0.0, 1.0]
}

func NewProbabilitySampler(rate float64) *ProbabilitySampler {
	return &ProbabilitySampler{rate: rate}
}

func (s *ProbabilitySampler) ShouldSample(traceID string) bool {
	return rand.Float64() < s.rate
}

// RateLimiterSampler 实现固定速率采样（如每秒 N 个）
type RateLimiterSampler struct {
	rate     int           // 每秒允许的最大采样数
	interval time.Duration // 计算间隔
	mu       sync.Mutex
	lastTime time.Time
	count    int
}

func NewRateLimiterSampler(ratePerSecond int) *RateLimiterSampler {
	return &RateLimiterSampler{
		rate:     ratePerSecond,
		interval: time.Second,
		lastTime: time.Now(),
	}
}

func (s *RateLimiterSampler) ShouldSample(traceID string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()
	if now.Sub(s.lastTime) >= s.interval {
		s.lastTime = now
		s.count = 0
	}

	if s.count < s.rate {
		s.count++
		return true
	}

	return false
}

// 根据traceID和请求URL决定是否采样
